require "lib"

local cjson = require("cjson")
local balancer = require "ngx.balancer"
local cache_data = require "cache_data"

local http_host = ngx.var.http_host
local cache_info = ngx.shared.cache_dict:get(GET_UPSTREAM_CACHE_KEY(http_host))
if not cache_info then
    local star_host = GET_MATCH_STAR_HOST(http_host)
    cache_info = ngx.shared.cache_dict:get(GET_UPSTREAM_CACHE_KEY(star_host))
end

if not cache_info then
    cache_info = ngx.shared.cache_dict:get(GET_UPSTREAM_CACHE_KEY("*"))
end

if not cache_info then
    ngx.log(ngx.ERR, "not set upstream")
    return ngx.exit(500)
end

local upstream_json, from_cache = cache_data.get_cache_to_json(GET_UPSTREAM_CACHE_KEY(http_host), cache_info)
if not upstream_json then
    ngx.log(ngx.ERR, "convert json failed")
    return ngx.exit(500)
end

local random_list = {}
local weight_list = {}
local all_weight = 0
local all_list = {}
for ip, data in pairs(upstream_json) do
    all_list[#all_list + 1] = ip
    local fail_times = tonumber(ngx.shared.cache_dict:get(ip .. "fail")) or 0
    if fail_times < (tonumber(data["fail"]) or 3) then
        random_list[#random_list + 1] = ip
        local weight = data["weight"] or 100
        table.insert(weight_list, weight)
        all_weight = all_weight + weight
    else
        local last_time = tonumber(ngx.shared.cache_dict:get(ip .. "last_time")) or 0
        if ngx.now() - last_time > (tonumber(data["fail_timeout"]) or 180) then
            ngx.shared.cache_dict:set(ip .. "fail", 0)
            random_list[#random_list + 1] = ip
            local weight = data["weight"] or 100
            table.insert(weight_list, weight)
            all_weight = all_weight + weight
        end
    end
end

local method = GET_BALANCER_METHOD()
local backend = nil

if #random_list > 0 then
    -- 随机权重
    if method == "weight" then
        local random_weight = math.random(0, all_weight)
        local now_weight = 0
        for idx, weight in ipairs(weight_list) do
            now_weight = now_weight + weight
            if random_weight <= now_weight then
                backend = random_list[idx]
                break
            end
        end
    -- Hash映射
    elseif method == "hash" then
        local port = ngx.var.server_port
        local remote_ip = GET_CLIENT_IP()
        local key = remote_ip..port
        local hash = ngx.crc32_long(key);
        hash = (hash % #random_list)
        if hash == 0 then
            hash = #random_list
        end
        backend = random_list[hash]
    end
end

if backend == nil then
    local host = random_list
    if #host == 0 then
        host = all_list
    end
    if #host == 0 then
        ngx.log(ngx.ERR, "failed get host ")
        return ngx.exit(500)
    end
    backend = host[math.random(1, #host)]
end

local new_port = 80
local idx = string.find(backend, ":")
if idx then
    new_port = tonumber(string.sub(backend, idx + 1, #backend)) or 80 
    backend = string.sub(backend, 1, idx - 1)
end
local ok, err = balancer.set_current_peer(backend, new_port)
if not ok then
    ngx.log(ngx.ERR, "failed to set the current peer: ", err)
    return ngx.exit(500)
end
ok, err = balancer.set_timeouts(3, 60, 60)